import pygame
import time
import random

WIDTH = 800
HEIGHT = 800

PADDING = 40
GRID_WIDTH = 240
GRID_HEIGHT = 240
GRID_PADDING = 20

CROSS = 1
CIRCLE = 2

CROSS_WIN = 1
CIRCLE_WIN = 2
DRAW = 3
NORMAL = 4

pygame.init()

screen = pygame.display.set_mode((WIDTH, HEIGHT))

class Board(object):
    def __init__(self):
        self.turn = CROSS
        self.game_over = False
        self.board = [
            [0, 0, 0],
            [0, 0, 0],
            [0, 0, 0],
        ]

        # 创建机器人
        self.ai = AI(self.board)

    def is_win(self, row, line):
        v = self.ai.has_won(row, line)
        if (v == CROSS_WIN):
            print("Cross Win")
        elif (v == CIRCLE_WIN):
            print("Circle Win")
        elif (v == DRAW):
            print("DRAW")

        if (v != NORMAL):
            self.game_over = True

    def click(self, pos):
        if (self.game_over):
            return

        x = pos[0]
        y = pos[1]

        if (x < PADDING or x > WIDTH - PADDING):
            return

        if (y < PADDING or y > HEIGHT - PADDING):
            return

        row = (y - 40) // GRID_HEIGHT
        line = (x - 40) // GRID_WIDTH

        if (self.board[row][line]):
            return

        self.board[row][line] = self.turn
        self.is_win(row, line)

        if (self.game_over):
            return

        # 让机器人自动下棋
        pos = self.ai.get_next(3 - self.turn)
        if (pos):
            self.is_win(pos[0], pos[1])
        else:
            print("平局")

    def draw_cross(self, row, line):
        center = (PADDING + line * GRID_WIDTH + GRID_WIDTH // 2,
            PADDING + row * GRID_HEIGHT + GRID_HEIGHT // 2)

        x1 = center[0] + GRID_WIDTH // 2 - GRID_PADDING
        y1 = center[1] - GRID_HEIGHT // 2 + GRID_PADDING
        
        x2 = center[0] - GRID_WIDTH // 2 + GRID_PADDING
        y2 = center[1] + GRID_HEIGHT // 2 - GRID_PADDING

        pygame.draw.line(screen, (0, 0, 0), (x1, y1), (x2, y2), 10)
        pygame.draw.line(screen, (0, 0, 0), (x2, y1), (x1, y2), 10)

    def draw_circle(self, row, line):
        center = (PADDING + line * GRID_WIDTH + GRID_WIDTH // 2,
            PADDING + row * GRID_HEIGHT + GRID_HEIGHT // 2)
        
        pygame.draw.circle(screen, (255, 0, 0), center, GRID_WIDTH // 2 - GRID_PADDING, 8)

    def draw(self, screen):
        screen.fill((192, 192, 192))
        # 画横线
        for i in range(4):
            pygame.draw.line(screen, (0, 0, 0), 
                (40, 40 + i * GRID_HEIGHT), (760, 40 + i * GRID_HEIGHT), 2)

        # 画竖线
        for i in range(4):
            pygame.draw.line(screen, (0, 0, 0), 
                (40 + i * GRID_WIDTH, 40), (40 + i * GRID_WIDTH, 760), 2)

        # 画棋子
        row = -1
        for l in self.board:
            row += 1
            line = -1
            for e in l:
                line += 1
                if e == CROSS:
                    self.draw_cross(row, line)
                elif e == CIRCLE:
                    self.draw_circle(row, line)

        pygame.display.update()


# Artificial Intelligence，人工智能
class AI(object):
    def __init__(self, board):
        self.board = board

    # row 行，line 列，是否组成三个连续的棋子
    def has_won(self, row, line):
        coin = self.board[row][line]

        # 垂直方向
        y = row - 1
        cnt = 1
        while (y >= 0 and self.board[y][line] == coin):
            cnt += 1
            y -= 1

        y = row + 1
        while (y < 3 and self.board[y][line] == coin):
            cnt += 1
            y += 1

        if (cnt == 3):
            print("垂直方向")
            if (coin == 1):
                return CROSS_WIN
            else:
                return CIRCLE_WIN

        # 水平方向
        x = line - 1
        cnt = 1
        while (x >= 0 and self.board[row][x] == coin):
            cnt += 1
            x -= 1

        x = line + 1
        while (x < 3 and self.board[row][x] == coin):
            cnt += 1
            x += 1

        if (cnt == 3):
            print("水平方向")
            if (coin == 1):
                return CROSS_WIN
            else:
                return CIRCLE_WIN

        # 斜向
        x = line - 1
        y = row - 1
        cnt = 1
        while (x >= 0 and y >= 0 and self.board[y][x] == coin):
            cnt += 1
            x -= 1
            y -= 1

        x = line + 1
        y = row + 1
        while (x < 3 and y < 3 and self.board[y][x] == coin):
            cnt += 1
            x += 1
            y += 1

        if (cnt == 3):
            print("斜方向")
            if (coin == 1):
                return CROSS_WIN
            else:
                return CIRCLE_WIN

        # 反斜向
        x = line - 1
        y = row + 1
        cnt = 1
        while (x >= 0 and y < 3 and self.board[y][x] == coin):
            cnt += 1
            x -= 1
            y += 1

        x = line + 1
        y = row - 1
        while (x < 3 and y >= 0 and self.board[y][x] == coin):
            cnt += 1
            x += 1
            y -= 1

        if (cnt == 3):
            print("反斜方向")
            if (coin == 1):
                return CROSS_WIN
            else:
                return CIRCLE_WIN

        cnt = 0
        for i in range(3):
            for j in range(3):
                if self.board[i][j] == 0:
                    cnt += 1

        if (cnt == 0):
            return DRAW
        else:
            return NORMAL

    def evaluate(self, row, line, turn):
        self.board[row][line] = turn

        # 如果有一方已经赢了
        win_score = self.has_won(row, line)
        if win_score == CROSS_WIN:
            return 100
        elif win_score == CICRCLE_WIN:
            return 0
        elif win_score == DRAW:
            return 50
                
        # 最大最小过程   
        value = 50
        for i in range(3):
            for j in range(3):
                if self.board[i][j] == 0:
                    score = self.evaluate(i, j, 3 - turn)
                    if turn == 1:
                        value = max(value, score)
                    else:
                        value = min(value, score)

        return value
                        

    # turn 取值为 1 或者 2
    def get_next(self, turn):
        lst = []
        # 统计所有的空白位置
        for i in range(3):
            for j in range(3):
                if self.board[i][j] == 0:
                    lst.append((i, j))

        if not lst:
            return None

        # 在所有的空白位置随机选一个
        y, x = random.choice(lst)
        self.board[y][x] = turn

        return y, x
        
        
board = Board()
board.draw(screen)

going = True

while going:
    for evt in pygame.event.get():
        if evt.type == pygame.QUIT:
            going = False
        elif evt.type == pygame.MOUSEBUTTONDOWN:
            board.click(evt.pos)
            board.draw(screen)

    time.sleep(0.02)
